"""
根据用户提出的问题，从向量数据库检索相关知识条目，发给大模型，让大模型给出一个理想的答案
"""
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import OpenAIEmbeddings, ChatOpenAI
from langchain_redis import RedisConfig, RedisVectorStore

from config.load_key import load_key

from operator import itemgetter

# 构建向量模型
embedding_model = OpenAIEmbeddings(
    model="Qwen/Qwen3-Embedding-0.6B",
    api_key=load_key("siliconflow_api_key"),
    base_url=load_key("siliconflow_base_url"),
)
# 向量数据持久化
config = RedisConfig(
    index_name="langchain_rag_test_index",
    redis_url=load_key("REDIS_URL"),
)
vector_store = RedisVectorStore(embedding_model, config)
retriever = vector_store.as_retriever()

# 检索问题
question = "主办方"
retriever_segments = retriever.invoke(question, k=5)
print(retriever_segments)

# 构建提示词模板
prompt_template = ChatPromptTemplate.from_messages([
    ("user", """你是一个答疑机器人，你的任务是根据下述给定的已知信息回答用户的问题。
    已知信息：{context}
    用户问题：{question}
    如果已知信息不包含用户问题的答案，或者已知信息不足以回答用户的问题，请直接回复”我无法回答您的问题“。
    请用中文回答用户问题""")
])
text = []
for segment in retriever_segments:
    text.append(segment.page_content)

prompt = prompt_template.invoke({"context": text, "question": question})
print(prompt.to_messages()[0].content)

llm = ChatOpenAI(
    model="Qwen/Qwen2.5-7B-Instruct",
    api_key=load_key("siliconflow_api_key"),
    base_url=load_key("siliconflow_base_url"),
)

response = llm.invoke(prompt)
print(response.content)


# 收集 document 的内容
def collect_documents(segments):
    text_array = []
    for item in segments:
        text_array.append(item.page_content)
    return text_array


# 构建链
chain = ({
    "context": itemgetter("question") | retriever | collect_documents,
    "question": itemgetter("question")
    }
    | prompt_template
    | llm
    | StrOutputParser()
)

response = chain.invoke({"question": question})
print(response)
